# -*- coding: utf-8 -*-
"""
Created on Mon Nov 18 13:00:38 2013

@author: wudong

File location :
'D:/MyCodeLib/MyAbaqusLib/'

usage :
import drmHelper as dh

Reload the updated module :
reload
import drmHelper as dh
reload(ah)

"""
from abaqusConstants import *
import regionToolset
import odbAccess as oa
import abaHelper as ah
import numpy as np
#
# ============================================================================
#
def openOdb(odbFilePath, readOnly=True):
	odb = oa.openOdb(odbFilePath, readOnly=readOnly)
	return odb
#
# ============================================================================
#
def getOdbNodeLabelFromCoord(odb, nodeCoord, tolerance=0.1):
	coord1 = nodeCoord
	if len(coord1)<3:
		coord1 = (nodeCoord[0], nodeCoord[1], 0)
	instance = odb.rootAssembly.instances.values()[0]
	allNodes = instance.nodes
	nodeLabel = -1
	(x1,y1,z1) = coord1
	for node in allNodes:
		(x2,y2,z2) = node.coordinates
		dist = pow(x1-x2,2)+pow(y1-y2,2)+pow(z1-z2,2)
		if dist<tolerance:
			nodeLabel = node.label
			break
	# assert (nodeLabel>=0), ('nodeLabel>=0, nodeCoord=(%f, %f, %f)'%(x1,y1,z1))
	return nodeLabel
#
# ============================================================================
#
def getOdbNodeLabelsFromNodeSet(odb, nodeSetName):
	instance = odb.rootAssembly.instances.values()[0]
	allNodes = instance.nodeSets[nodeSetName].nodes
	nodeLabels = []
	for node in allNodes:
		nodeLabels.append(node.label)
	return nodeLabels
#
# ============================================================================
#
def createOdbNodeSetFromLabels(odb, nodeSetName, nodeLabels):
	instance = odb.rootAssembly.instances.values()[0]
	nodeSet = instance.NodeSetFromNodeLabels(name=nodeSetName, nodeLabels=tuple(nodeLabels))
	return nodeSet
#
# ============================================================================
#
def getOdbElemLabelsFromElemSet(odb, elemSetName):
	instance = odb.rootAssembly.instances.values()[0]
	allElems = instance.elementSets[elemSetName].elements
	elemLabels = []
	for node in allElems:
		elemLabels.append(node.label)
	return elemLabels
#
# ============================================================================
#
def createOdbElemSetFromLabels(odb, elemSetName, elemLabels):
	instance = odb.rootAssembly.instances.values()[0]
	elemSet = instance.ElementSetFromElementLabels(name=elemSetName, elementLabels=tuple(elemLabels))
	return elemSet
#
# ============================================================================
#
def getOdbNodeHistoryOutputData(odb, stepName, nodeLabel, outputVarName):
	step = odb.steps[stepName]
	instance = odb.rootAssembly.instances.values()[0]
	instanceName = instance.name
	regionName = 'Node %s.%d' % (instanceName, nodeLabel)
	data = step.historyRegions[regionName].historyOutputs[outputVarName].data
	data = ah.tableToArray(data)
	return data
#
# ============================================================================
#
def getFrameFieldOutputData(frame,varName,region):
	return frame.fieldOutputs[varName].getSubset(region=region)
#
# ============================================================================
#
def getElemFieldOutputHistData(step,outputVarNames,elemRegion=MODEL):
	elemCount = len(elemRegion.elements)
	frameCount = len(step.frames)
	varCount = len(outputVarNames)
	elemLabels = np.zeros(shape=(elemCount,varCount))
	frameInfo = np.zeros(shape=(frameCount,2))
	outputData = np.zeros(shape=(elemCount,frameCount,varCount))
	for j in range(frameCount):
		frame = step.frames[j]
		frameNum = frame.incrementNumber
		frameTime = frame.frameValue
		frameInfo[j,0] = frameNum
		frameInfo[j,1] = frameTime
		for k in range(varCount):
			varName = outputVarNames[k]
			varData = getFrameFieldOutputData(frame,varName,elemRegion)
			for i in range(elemCount):
				v = varData.values[i]
				elemLabels[i,k] = v.elementLabel
				outputData[i,j,k] = v.data
	return (elemLabels,frameInfo,outputData)
#
# ============================================================================
#
def getNodeSetFieldOutputHistData(odb,stepName,nodeSetName,outputVarNames):
	instance = odb.rootAssembly.instances.values()[0]
	nodeRegion = instance.nodeSets[nodeSetName]
	nodeCount = len(nodeRegion.nodes)
	step = odb.steps[stepName]
	frameCount = len(step.frames)
	varCount = len(outputVarNames)
	nodeLabels = np.zeros(shape=(nodeCount,varCount))
	frameInfo = np.zeros(shape=(frameCount,2))
	outputData = np.zeros(shape=(nodeCount,frameCount,varCount))
	for j in range(frameCount):
		frame = step.frames[j]
		frameNum = frame.incrementNumber
		frameTime = frame.frameValue
		frameInfo[j,0] = frameNum
		frameInfo[j,1] = frameTime
		for k in range(varCount):
			varName = outputVarNames[k]
			varData = getFrameFieldOutputData(frame,varName,nodeRegion)
			for i in range(nodeCount):
				v = varData.values[i]
				nodeLabels[i,k] = v.nodeLabel
				outputData[i,j,k] = v.data
	return (nodeLabels,frameInfo,outputData)
#
# ============================================================================
#
from abaqusConstants import *
import regionToolset
import odbAccess as oa
import abaHelper as ah
#
class AbaqusOdb:
	def __init__(self, odbFilePath):
		self.odbFilePath = odbFilePath
		pass
	
	def openOdb(self, readOnly=True):
		self.odb = oa.openOdb(self.odbFilePath, readOnly=readOnly)
		self.instance = self.odb.rootAssembly.instances.values()[0]
		pass
	
	def closeOdb(self):
		self.odb.close()

	def getNodeLabelsFromNodeSet(self, nodeSetName):
		assert nodeSetName.upper() in self.instance.nodeSets.keys()
		allNodes = self.instance.nodeSets[nodeSetName.upper()].nodes
		nodeLabels = []
		for node in allNodes:
			nodeLabels.append(node.label)
		return nodeLabels
	
	def getNodeCoordinateFromNodeLabel(self, nodeLabel):
		allNodes = self.instance.nodes
		nodeCoord = None
		for node in allNodes:
			if node.label==nodeLabel:
				nodeCoord = node.coordinates
		assert nodeCoord!=None
		return nodeCoord
	
	def getNodeCoordinatesFromNodeSet(self, nodeSetName):
		assert nodeSetName.upper() in self.instance.nodeSets.keys()
		allNodes = self.instance.nodeSets[nodeSetName.upper()].nodes
		nodeCoords = []
		for node in allNodes:
			nodeCoords.append(node.coordinates)
		return nodeCoords

	def getNodeLabelFromNodeCoordinate(self, nodeCoord, tolerance=0.1):
		assert (type(nodeCoord) is tuple) and (len(nodeCoord)>=2)
		coord1 = nodeCoord
		if len(coord1)<3:
			coord1 = (nodeCoord[0], nodeCoord[1], 0)
		allNodes = self.instance.nodes
		nodeLabel = -1
		(x1,y1,z1) = coord1
		for node in allNodes:
			(x2,y2,z2) = node.coordinates
			dist = pow(x1-x2,2)+pow(y1-y2,2)+pow(z1-z2,2)
			if dist<tolerance:
				nodeLabel = node.label
				break
		assert (nodeLabel>=0), ('node not found, nodeCoord=(%f, %f, %f)' % (x1,y1,z1))
		return nodeLabel

	def getNodeHistoryOutput(self, stepName, nodeLabel, outputVarName):
		assert stepName in self.odb.steps.keys(), 'stepName=%s' % (stepName,)
		step = self.odb.steps[stepName]
		instanceName = self.instance.name
		regionName = 'Node %s.%d' % (instanceName, nodeLabel)
		data = step.historyRegions[regionName].historyOutputs[outputVarName].data
		data = ah.tableToArray(data)
		return data

	def getElemSetFieldOutputHistoryData(self, stepName, elemSetName, outputVarName):
		assert elemSetName.upper() in self.instance.elementSets.keys(), 'elem set not found : %s' % elemSetName
		assert stepName.upper() in self.odb.steps.keys()
		region = self.instance.elementSets[elemSetName]
		step = self.odb.steps[stepName]
		#
		elemCount = len(region.elements)
		frameCount = len(step.frames)
		elemLabels = np.zeros(shape=(elemCount,))
		frameTimes = np.zeros(shape=(frameCount,))
		outputData = np.zeros(shape=(elemCount,frameCount))
		for j in range(frameCount):
			frame = step.frames[j]
			frameTime = frame.frameValue
			frameTimes[j] = frameTime
			varData = frame.fieldOutputs[outputVarName].getSubset(region=region)
			for i in range(elemCount):
				v = varData.values[i]
				elemLabels[i] = v.elementLabel
				outputData[i,j] = v.data
		return (elemLabels, frameTimes, outputData)

	def getElemFieldOutputHistoryData(self, stepName, elemLabel, outputVarName, varComponentCount, timePointData):
		assert stepName in self.odb.steps.keys()
		timePointCount = timePointData.shape[0]
		elemData = np.zeros(shape=(timePointCount, varComponentCount))
		elemSetName = '_ESET_%d' % (elemLabel,)
		elemSet = None
		if elemSetName in self.instance.elementSets.keys():
			elemSet = self.instance.elementSets[elemSetName]
		else:
			elemSet = self.instance.ElementSetFromElementLabels(name=elemSetName,elementLabels=(elemLabel,))
		#
		assert len(elemSet.elements)>0
		elemType = elemSet.elements[0].type
		step = self.odb.steps[stepName]
		EPS = 1.0e-5
		i = 0
		for frame in step.frames:
			frameTime = frame.frameValue
			if abs(frameTime-timePointData[i])>EPS:
				continue
			#
			assert outputVarName.upper() in frame.fieldOutputs.keys()
			varData = frame.fieldOutputs[outputVarName].getSubset(region=elemSet, position=INTEGRATION_POINT, elementType=elemType)
			elemData[i, :] = np.array(varData.values[0].data)
			i = i+1
		#
		assert i>=timePointCount, 'i=%d' % (i,)
		return elemData

	def getElemSetVariable(self, stepName, elemSetName, elemVarName, frameNum=0):
		assert stepName in self.odb.steps.keys()
		assert elemSetName in self.instance.elementSets.keys()
		elemSet = self.instance.elementSets[elemSetName]
		elemCount = len(elemSet.elements)
		assert elemCount>0
		elemData = np.zeros(shape=(elemCount, 2))
		#
		elemType = elemSet.elements[0].type
		step = self.odb.steps[stepName]
		frame = step.frames[frameNum]
		assert elemVarName.upper() in frame.fieldOutputs.keys()
		varData = frame.fieldOutputs[elemVarName].getSubset(region=elemSet)
		for i in range(elemCount):
			v = varData.values[i]
			elemData[i,0] = v.elementLabel	# elemNum
			elemData[i,1] = v.data			# variable
		#
		return elemData

